"""
cron: */5 * * * *
new Env('股票交易信号专业版');
"""
import requests
import os
import mysql.connector
import json
import time
from datetime import datetime
import numpy as np
import talib

class StockSignalPro:
    def __init__(self):
        self.connection = mysql.connector.connect(
            host=os.environ.get('HOST'),
            user=os.environ.get('USER'),
            password=os.environ.get('PASSWORD'),
            database=os.environ.get('DATABASE')
        )
        self.cursor = self.connection.cursor()
        self.create_tables()
        
    def create_tables(self):
        """创建必要的数据表"""
        # 股票信号表
        self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS ningm_stock_signals (
                id INT AUTO_INCREMENT PRIMARY KEY,
                stock_code VARCHAR(10),
                stock_name VARCHAR(100),
                signal_type VARCHAR(20),
                signal_strength INT,
                current_price DECIMAL(10,2),
                target_price DECIMAL(10,2),
                stop_loss DECIMAL(10,2),
                indicators TEXT,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                INDEX idx_stock_code (stock_code),
                INDEX idx_signal_type (signal_type),
                INDEX idx_create_time (create_time)
            )
        """)
        
        # 技术指标表
        self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS ningm_technical_indicators (
                id INT AUTO_INCREMENT PRIMARY KEY,
                stock_code VARCHAR(10),
                ma5 DECIMAL(10,2),
                ma10 DECIMAL(10,2),
                ma20 DECIMAL(10,2),
                rsi DECIMAL(10,2),
                kdj_k DECIMAL(10,2),
                kdj_d DECIMAL(10,2),
                kdj_j DECIMAL(10,2),
                macd DECIMAL(10,2),
                boll_up DECIMAL(10,2),
                boll_mid DECIMAL(10,2),
                boll_low DECIMAL(10,2),
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                INDEX idx_stock_code (stock_code),
                INDEX idx_create_time (create_time)
            )
        """)
        
        # 策略配置表
        self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS ningm_strategy_config (
                id INT AUTO_INCREMENT PRIMARY KEY,
                strategy_name VARCHAR(50),
                indicators TEXT,
                conditions TEXT,
                weights TEXT,
                user_id INT,
                is_active BOOLEAN DEFAULT TRUE,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                INDEX idx_user_id (user_id),
                INDEX idx_strategy_name (strategy_name)
            )
        """)

    def get_stock_data(self, stock_code):
        """获取股票数据"""
        try:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            
            # 使用新浪股票API
            url = f'http://hq.sinajs.cn/list={stock_code}'
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                data = response.text.split('=')[1].strip('"').split(',')
                return {
                    'code': stock_code,
                    'name': data[0],
                    'open': float(data[1]),
                    'close': float(data[3]),
                    'high': float(data[4]),
                    'low': float(data[5]),
                    'volume': float(data[8]),
                    'amount': float(data[9])
                }
                
        except Exception as e:
            print(f"获取股票数据失败: {str(e)}")
            return None

    def calculate_indicators(self, stock_code, prices):
        """计算技术指标"""
        try:
            close_prices = np.array(prices)
            
            # 计算移动平均线
            ma5 = talib.MA(close_prices, timeperiod=5)[-1]
            ma10 = talib.MA(close_prices, timeperiod=10)[-1]
            ma20 = talib.MA(close_prices, timeperiod=20)[-1]
            
            # 计算RSI
            rsi = talib.RSI(close_prices, timeperiod=14)[-1]
            
            # 计算KDJ
            k, d = talib.STOCH(high=close_prices, 
                              low=close_prices, 
                              close=close_prices,
                              fastk_period=9,
                              slowk_period=3,
                              slowk_matype=0,
                              slowd_period=3,
                              slowd_matype=0)
            j = 3 * k[-1] - 2 * d[-1]
            
            # 计算MACD
            macd, signal, hist = talib.MACD(close_prices)
            
            # 计算布林带
            upper, middle, lower = talib.BBANDS(close_prices, 
                                              timeperiod=20,
                                              nbdevup=2,
                                              nbdevdn=2,
                                              matype=0)
            
            indicators = {
                'stock_code': stock_code,
                'ma5': ma5,
                'ma10': ma10,
                'ma20': ma20,
                'rsi': rsi,
                'kdj_k': k[-1],
                'kdj_d': d[-1],
                'kdj_j': j,
                'macd': macd[-1],
                'boll_up': upper[-1],
                'boll_mid': middle[-1],
                'boll_low': lower[-1]
            }
            
            self.save_indicators(indicators)
            return indicators
            
        except Exception as e:
            print(f"计算技术指标失败: {str(e)}")
            return None

    def generate_signals(self, stock_data, indicators):
        """生成交易信号"""
        try:
            signals = []
            current_price = stock_data['close']
            
            # 均线多头排列
            if indicators['ma5'] > indicators['ma10'] > indicators['ma20']:
                signals.append({
                    'type': 'BUY',
                    'strength': 2,
                    'reason': '均线多头排列'
                })
            
            # RSI超买超卖
            if indicators['rsi'] < 30:
                signals.append({
                    'type': 'BUY',
                    'strength': 1,
                    'reason': 'RSI超卖'
                })
            elif indicators['rsi'] > 70:
                signals.append({
                    'type': 'SELL',
                    'strength': 1,
                    'reason': 'RSI超买'
                })
            
            # KDJ金叉死叉
            if indicators['kdj_j'] > indicators['kdj_d'] and \
               indicators['kdj_d'] > indicators['kdj_k']:
                signals.append({
                    'type': 'BUY',
                    'strength': 2,
                    'reason': 'KDJ金叉'
                })
            elif indicators['kdj_j'] < indicators['kdj_d'] and \
                 indicators['kdj_d'] < indicators['kdj_k']:
                signals.append({
                    'type': 'SELL',
                    'strength': 2,
                    'reason': 'KDJ死叉'
                })
            
            # MACD金叉死叉
            if indicators['macd'] > 0:
                signals.append({
                    'type': 'BUY',
                    'strength': 3,
                    'reason': 'MACD金叉'
                })
            elif indicators['macd'] < 0:
                signals.append({
                    'type': 'SELL',
                    'strength': 3,
                    'reason': 'MACD死叉'
                })
            
            # 布林带突破
            if current_price > indicators['boll_up']:
                signals.append({
                    'type': 'SELL',
                    'strength': 1,
                    'reason': '突破布林带上轨'
                })
            elif current_price < indicators['boll_low']:
                signals.append({
                    'type': 'BUY',
                    'strength': 1,
                    'reason': '突破布林带下轨'
                })
            
            return signals
            
        except Exception as e:
            print(f"生成交易信号失败: {str(e)}")
            return []

    def save_signals(self, stock_data, signals):
        """保存交易信号"""
        try:
            for signal in signals:
                insert_query = """
                    INSERT INTO ningm_stock_signals 
                    (stock_code, stock_name, signal_type, signal_strength,
                     current_price, target_price, stop_loss, indicators) 
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """
                
                # 计算目标价和止损价
                current_price = stock_data['close']
                if signal['type'] == 'BUY':
                    target_price = current_price * 1.1  # 上涨10%
                    stop_loss = current_price * 0.95    # 下跌5%
                else:
                    target_price = current_price * 0.9  # 下跌10%
                    stop_loss = current_price * 1.05    # 上涨5%
                
                insert_values = (
                    stock_data['code'],
                    stock_data['name'],
                    signal['type'],
                    signal['strength'],
                    current_price,
                    target_price,
                    stop_loss,
                    json.dumps({'reason': signal['reason']})
                )
                self.cursor.execute(insert_query, insert_values)
                
            self.connection.commit()
            
        except Exception as e:
            print(f"保存交易信号失败: {str(e)}")
            self.connection.rollback()

    def save_indicators(self, indicators):
        """保存技术指标"""
        try:
            insert_query = """
                INSERT INTO ningm_technical_indicators 
                (stock_code, ma5, ma10, ma20, rsi, kdj_k, kdj_d, kdj_j,
                 macd, boll_up, boll_mid, boll_low) 
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            insert_values = (
                indicators['stock_code'],
                indicators['ma5'],
                indicators['ma10'],
                indicators['ma20'],
                indicators['rsi'],
                indicators['kdj_k'],
                indicators['kdj_d'],
                indicators['kdj_j'],
                indicators['macd'],
                indicators['boll_up'],
                indicators['boll_mid'],
                indicators['boll_low']
            )
            self.cursor.execute(insert_query, insert_values)
            self.connection.commit()
            
        except Exception as e:
            print(f"保存技术指标失败: {str(e)}")
            self.connection.rollback()

    def run(self):
        """运行信号系统"""
        try:
            # 股票代码列表（示例）
            stock_codes = ['sh000001', 'sz399001', 'sh600519', 'sz000858']
            
            for stock_code in stock_codes:
                # 1. 获取股票数据
                stock_data = self.get_stock_data(stock_code)
                if not stock_data:
                    continue
                
                # 2. 计算技术指标
                indicators = self.calculate_indicators(stock_code, 
                                                    [stock_data['close']])
                if not indicators:
                    continue
                
                # 3. 生成交易信号
                signals = self.generate_signals(stock_data, indicators)
                
                # 4. 保存交易信号
                if signals:
                    self.save_signals(stock_data, signals)
                
                time.sleep(1)  # 避免请求过快
                
        except Exception as e:
            print(f"系统运行失败: {str(e)}")
        finally:
            self.cursor.close()
            self.connection.close()

if __name__ == "__main__":
    signal_system = StockSignalPro()
    signal_system.run() 